#!/usr/bin/env python
# coding=UTF-8
"""
@Author: linna
@LastEditors: linna
@Description:
@Date: 2020-05-08 14:50:02
@LastEditTime: 2020-05-09 09:11:02
"""
import numpy as np
import torch
from PIL import Image

from .makes import ImageTransfer


class SYSTEM:
    def __init__(self, root_dir, meta_file, config=None, transform_type='val'):
        """
        @description:
        @param {
            decoder_type = ['pil', 'opencv', 'ffmpeg']
            resize_type = ['pil-bilinear', 'pil-nearest', 'pil-box',
            'pil-hamming', 'pil-cubic', 'pil-lanczos',
            'opencv-nearest', 'opencv-bilinear',
            'opencv-area', 'opencv-cubic', 'opencv-lanczos',]                
        }
        @return: None
        """
        self.root_dir = root_dir
        self.meta_file = meta_file
        self.transform_type = transform_type
        self.decoder_type = 'pil'
        self.resize_type = 'pil-bilinear'
        self._parse_params(config)
        self.cur_id = 0
        self.system_gen = ImageTransfer(
            root_dir=self.root_dir,
            meta_file=self.meta_file,
            save_root='',
            decoder_type=self.decoder_type,
            resize_type=self.resize_type,
            transform_type=self.transform_type,
        )
        print(f"decoder_type is {self.decoder_type}, resize_type is {self.resize_type}")

    def _parse_params(self, config):
        self.decoder_type = config.get("decoder_type", "pil")
        self.resize_type = config.get("resize_type", 'pil-bilinear')

    def generate(self, xs=None, ys=None):
        """
        @description:
        @param {
            xs:
            ys:
            device:
        }
        @return:
        """
        copy_xs = xs.permute(0, 2, 3, 1).numpy()
        cor_xs = []

        for _i, x in enumerate(copy_xs):
            # x = (x * 255).astype(np.uint8)

            cor_x, label = self.system_gen.getimage(idx=self.cur_id)
            cor_xs.append(torch.from_numpy(cor_x / 255.0))
            self.cur_id += 1
        cor_xs = torch.stack(cor_xs, 0).permute(0, 3, 1, 2)
        cor_xs = torch.tensor(cor_xs, dtype=torch.float32)

        return cor_xs
